/*
  Copyright 2005 - Daniel Lowd

  This file is part of the nbe code available at http://www.cs.washington.edu/ai/nbe/. There are some modifications to 
  integrate the original code of Daniel Lowd into the eapmlib library. The NBE method is explained in:

  Daniel Lowd and Pedro Domingos. 'Naive Bayes Models for Probability Estimation'. 
  Proceedings of the Twenty-Second International Conference on Machine Learning (ICML), 2005. 
  Bonn, Germany: ACM Press.
*/
/*!
\file AbstractionSet.h
\brief AbstractionSet declaration. This defines the probabilistic model.
\author Daniel Lowd

  This class defines a set of abstractions, it is a set of clusters. This set of clusters is the Naive Bayes Model.
*/
#ifndef __ABSTRACTIONSET_H
#define __ABSTRACTIONSET_H

#include "EvolutiveLib.h"

#include "GrowArray.h"
#include "VarSchema.h"
#include "VarSet.h"
#include "Prob.h"
#include <stdlib.h>
#include <iostream>

// HACK (for margtest)
#include <vector>
#include <list>
#include "GenDistribution.h"
#include "Distribution.h"
#include "Abstraction.h"

using namespace std;

namespace NBE {

typedef struct 
{
	double weight;
	Abstraction* abs;
} absWeight;

class AbstractionSet
{
protected:
    GrowArray<Abstraction*> abstractions;
    GrowArray<double> weights;
    int numAbs;
    VarSchema schema;
    int numVars;

public:
    AbstractionSet() :numAbs(0), numVars(0) {};
    AbstractionSet(const VarSchema& s, int initialSize = 16);
    virtual ~AbstractionSet();

    // NOTE: AbstractionSet takes ownership of a.
    void addAbstraction(Abstraction* a);
    void addNullAbstraction();
    void removeLastAbstraction();

    int findAbstraction(Abstraction* abs) const;
    virtual VarSet getSample() const;
    virtual VarSet getSample(int a) const = 0;

    int getNumAbstractions() const { return numAbs; }
    double getWeight(int i) const { return weights[i]; }
    const Abstraction& getAbstraction(int i) const {return *abstractions[i];}

    double getLogProb(const VarSet& x) const { return getLogLikelihood(x); }
    double getLogLikelihood(const VarSet& x) const;
    double getLogLikelihood(const VarSet* obs, int numObs) const;
    virtual Prob getProb(const VarSet& x, int a) const;
    virtual double getProb(int varIndex, double varVal, int a) const;

    // HACK (for margtest) -- Implemented only in AbsVarAbstractionSet
    virtual vector<GenDistribution*> createMarginals(const VarSet& x) const = 0;
    // HACK (for jointtest) -- Implemented only in AbsVarAbstractionSet
    virtual AbstractionSet* createJoint(list<int> queryVars, 
            const VarSet& evidence) const = 0;
    virtual int getRange(int i) const { return schema.getRange(i); }

    // Weight all abstractions equally.
    void resetWeights();

    // Take a single step towards learning the proper abstraction weights
    // using the EM algorithm.
    double EMStep(VarSet* obs, int numObs);

    // Keep only a set fraction of abstractions (by total weight),
    // or those with minimum weight. 
    void prune(double weightFraction);
    void threshold(double threshold);

    // EM hooks for use in subclasses.
    virtual void initParameters(VarSet* obs, int numObs) {}
    virtual void extendParameters() {}
    virtual void assignObservation(const VarSet& obs, int a, double p) {}
    virtual void updateParameters() {}
    virtual void saveParameters() {}
    virtual void restoreParameters() {}
    virtual void resetParameters() {}


    // I/O
    istream& readAbstractions(istream& in);
    ostream& writeAbstractions(ostream& out) const;
    istream& readSchema(istream& in);
    ostream& writeSchema(ostream& out) const;
    istream& readWeights(istream& in);
    ostream& writeWeights(ostream& out) const;
    virtual istream& readParameters(istream& in) { return in; }
    virtual ostream& writeParameters(ostream& out) const { return out; }
    
    friend istream& operator>>(istream& in, AbstractionSet& as);
    friend ostream& operator<<(ostream& out, const AbstractionSet& as);

private:
	//! Xevi: Change the random rutine
    //double frand() const { return random()/(double)RAND_MAX; }
	double frand() const { return RAND_VALUE(); }

    // Disallow copy constructor and assignment op
    AbstractionSet(const AbstractionSet& other);
    AbstractionSet& operator=(const AbstractionSet& other);
};

istream& operator>>(istream& in, AbstractionSet& model);
ostream& operator<<(ostream& out, const AbstractionSet& model);

}

#endif // __ABSTRACTIONSET_H
